from core.algorithm.modules.tracking.tracker.ocsort_tracker import OCSORTTracker
from core.algorithm.modules.tracking.tracker.botsort_tracker import BOTSORTTracker
import yaml
import numpy as np


class Track():
    """追踪模块, 将追踪算法整合为处理模块
    """
    def __init__(self, keys,config_path) -> None:
        self.keys= keys
        with open(config_path) as f:
            track_cfg = yaml.safe_load(f)
        track_type = track_cfg["type"]
        track_cfg = track_cfg[track_type]
        if track_type == "OCSORTTracker":
            self.tracker = OCSORTTracker(**track_cfg)
        else:
            raise Exception("|不支持的追踪模型类型|")

    def __call__(self, data):
        det_res = data[self.keys["in"]]
        input = det_res["boxes"]
        track_res = self.tracker.update(input) # [x0, y0, x1, y1, score, cls, id]
        det_res["boxes"] = track_res
        det_res["boxes_num"] = np.array([track_res.shape[0]])
        data[self.keys["out"]] = det_res
        return data
    
class MMTrack():
    """追踪模块, 将追踪算法整合为处理模块
    """
    def __init__(self, keys,
                 MMTrack_batch_size = 1, 
                 MMTrack_cfg = "core/algorithm/modules/tracking/tracker_config.yml",
                 ) -> None:
        self.keys= keys
        self.camera_move = False
        with open(MMTrack_cfg) as f:
            track_cfg = yaml.safe_load(f)
        track_type = track_cfg["type"]
        track_cfg = track_cfg[track_type]
        if track_type == "OCSORTTracker":
            self.tracker = [OCSORTTracker(**track_cfg) for _ in range(MMTrack_batch_size)]
        elif track_type == "BOTSORTTracker":
            self.tracker = [BOTSORTTracker(**track_cfg) for _ in range(MMTrack_batch_size)]
            self.camera_move = track_cfg["camera_motion"]
        else:
            raise Exception("|不支持的追踪模型类型|")

    def __call__(self, data):
        img_batch = data[self.keys["in_data"]]
        det_res = data[self.keys["in_det"]]
        batch_boxes = det_res["boxes"]
        res_boxes= []
        track_ids = []
        for img_data, boxes, tracker in zip(img_batch, batch_boxes, self.tracker):
            if self.camera_move:
                boxes = tracker.update(boxes, img_data)
            else:
                boxes = tracker.update(boxes)
            res_boxes.append(boxes[:,:-1])
            track_ids.append(boxes[:,-1])
        data[self.keys["out_track"]] = {
            "boxes": res_boxes,
            "track_ids": track_ids,
        }
        return data